// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::{
    schema::{Name, Schema},
    types::Value,
    util::{zig_i32, zig_i64},
};
use std::{collections::HashMap, convert::TryInto};

/// Encode a `Value` into avro format.
///
/// **NOTE** This will not perform schema validation. The value is assumed to
/// be valid with regards to the schema. Schema are needed only to guide the
/// encoding for complex type values.
pub fn encode(value: &Value, schema: &Schema, buffer: &mut Vec<u8>) {
    encode_ref(value, schema, buffer)
}

fn encode_bytes<B: AsRef<[u8]> + ?Sized>(s: &B, buffer: &mut Vec<u8>) {
    let bytes = s.as_ref();
    encode(&Value::Long(bytes.len() as i64), &Schema::Long, buffer);
    buffer.extend_from_slice(bytes);
}

fn encode_long(i: i64, buffer: &mut Vec<u8>) {
    zig_i64(i, buffer)
}

fn encode_int(i: i32, buffer: &mut Vec<u8>) {
    zig_i32(i, buffer)
}

/// Encode a `Value` into avro format.
///
/// **NOTE** This will not perform schema validation. The value is assumed to
/// be valid with regards to the schema. Schema are needed only to guide the
/// encoding for complex type values.
pub fn encode_ref(value: &Value, schema: &Schema, buffer: &mut Vec<u8>) {
    fn encode_ref0(
        value: &Value,
        schema: &Schema,
        buffer: &mut Vec<u8>,
        schemas_by_name: &mut HashMap<Name, Schema>,
    ) {
        match &schema {
            Schema::Ref { ref name } => {
                let resolved = schemas_by_name.get(name).unwrap();
                return encode_ref0(value, resolved, buffer, &mut schemas_by_name.clone());
            }
            Schema::Record { ref name, .. }
            | Schema::Enum { ref name, .. }
            | Schema::Fixed { ref name, .. } => {
                schemas_by_name.insert(name.clone(), schema.clone());
            }
            _ => (),
        }

        match value {
            Value::Null => (),
            Value::Boolean(b) => buffer.push(if *b { 1u8 } else { 0u8 }),
            // Pattern | Pattern here to signify that these _must_ have the same encoding.
            Value::Int(i) | Value::Date(i) | Value::TimeMillis(i) => encode_int(*i, buffer),
            Value::Long(i)
            | Value::TimestampMillis(i)
            | Value::TimestampMicros(i)
            | Value::TimeMicros(i) => encode_long(*i, buffer),
            Value::Float(x) => buffer.extend_from_slice(&x.to_le_bytes()),
            Value::Double(x) => buffer.extend_from_slice(&x.to_le_bytes()),
            Value::Decimal(decimal) => match schema {
                Schema::Decimal { inner, .. } => match *inner.clone() {
                    Schema::Fixed { size, .. } => {
                        let bytes = decimal.to_sign_extended_bytes_with_len(size).unwrap();
                        let num_bytes = bytes.len();
                        if num_bytes != size {
                            panic!(
                                "signed decimal bytes length {} not equal to fixed schema size {}",
                                num_bytes, size
                            );
                        }
                        encode(&Value::Fixed(size, bytes), inner, buffer)
                    }
                    Schema::Bytes => {
                        encode(&Value::Bytes(decimal.try_into().unwrap()), inner, buffer)
                    }
                    _ => panic!("invalid inner type for decimal: {:?}", inner),
                },
                _ => panic!("invalid schema type for decimal: {:?}", schema),
            },
            &Value::Duration(duration) => {
                let slice: [u8; 12] = duration.into();
                buffer.extend_from_slice(&slice);
            }
            Value::Uuid(uuid) => encode_bytes(&uuid.to_string(), buffer),
            Value::Bytes(bytes) => match *schema {
                Schema::Bytes => encode_bytes(bytes, buffer),
                Schema::Fixed { .. } => buffer.extend(bytes),
                _ => error!("invalid schema type for bytes: {:?}", schema),
            },
            Value::String(s) => match *schema {
                Schema::String => {
                    encode_bytes(s, buffer);
                }
                Schema::Enum { ref symbols, .. } => {
                    if let Some(index) = symbols.iter().position(|item| item == s) {
                        encode_int(index as i32, buffer);
                    }
                }
                _ => error!("invalid schema type for String: {:?}", schema),
            },
            Value::Fixed(_, bytes) => buffer.extend(bytes),
            Value::Enum(i, _) => encode_int(*i as i32, buffer),
            Value::Union(idx, item) => {
                if let Schema::Union(ref inner) = *schema {
                    inner.schemas.iter().for_each(|s| match s {
                        Schema::Record { name, .. }
                        | Schema::Enum { name, .. }
                        | Schema::Fixed { name, .. } => {
                            schemas_by_name.insert(name.clone(), s.clone());
                        }
                        _ => (),
                    });

                    let inner_schema = inner
                        .schemas
                        .get(*idx as usize)
                        .expect("Invalid Union validation occurred");
                    encode_long(*idx as i64, buffer);
                    encode_ref0(&*item, inner_schema, buffer, schemas_by_name);
                } else {
                    error!("invalid schema type for Union: {:?}", schema);
                }
            }
            Value::Array(items) => {
                if let Schema::Array(ref inner) = *schema {
                    if !items.is_empty() {
                        encode_long(items.len() as i64, buffer);
                        for item in items.iter() {
                            encode_ref0(item, inner, buffer, schemas_by_name);
                        }
                    }
                    buffer.push(0u8);
                } else {
                    error!("invalid schema type for Array: {:?}", schema);
                }
            }
            Value::Map(items) => {
                if let Schema::Map(ref inner) = *schema {
                    if !items.is_empty() {
                        encode_long(items.len() as i64, buffer);
                        for (key, value) in items {
                            encode_bytes(key, buffer);
                            encode_ref0(value, inner, buffer, schemas_by_name);
                        }
                    }
                    buffer.push(0u8);
                } else {
                    error!("invalid schema type for Map: {:?}", schema);
                }
            }
            Value::Record(fields) => {
                if let Schema::Record {
                    fields: ref schema_fields,
                    ..
                } = *schema
                {
                    for (i, &(_, ref value)) in fields.iter().enumerate() {
                        encode_ref0(value, &schema_fields[i].schema, buffer, schemas_by_name);
                    }
                }
            }
        }
    }

    let mut schemas_by_name = HashMap::new();
    encode_ref0(value, schema, buffer, &mut schemas_by_name)
}

pub fn encode_to_vec(value: &Value, schema: &Schema) -> Vec<u8> {
    let mut buffer = Vec::new();
    encode(value, schema, &mut buffer);
    buffer
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn test_encode_empty_array() {
        let mut buf = Vec::new();
        let empty: Vec<Value> = Vec::new();
        encode(
            &Value::Array(empty),
            &Schema::Array(Box::new(Schema::Int)),
            &mut buf,
        );
        assert_eq!(vec![0u8], buf);
    }

    #[test]
    fn test_encode_empty_map() {
        let mut buf = Vec::new();
        let empty: HashMap<String, Value> = HashMap::new();
        encode(
            &Value::Map(empty),
            &Schema::Map(Box::new(Schema::Int)),
            &mut buf,
        );
        assert_eq!(vec![0u8], buf);
    }

    #[test]
    fn test_avro_3433_recursive_definition_encode_record() {
        let mut buf = Vec::new();
        let schema = Schema::parse_str(
            r#"
        {
            "type":"record",
            "name":"TestStruct",
            "fields": [
                {
                    "name":"a",
                    "type":{
                        "type":"record",
                        "name": "Inner",
                        "fields": [ {
                            "name":"z",
                            "type":"int"
                        }]
                    }
                },
                {
                    "name":"b",
                    "type":"Inner"
                }
            ]
        }"#,
        )
        .unwrap();

        let inner_value1 = Value::Record(vec![("z".into(), Value::Int(3))]);
        let inner_value2 = Value::Record(vec![("z".into(), Value::Int(6))]);
        let outer_value =
            Value::Record(vec![("a".into(), inner_value1), ("b".into(), inner_value2)]);
        encode(&outer_value, &schema, &mut buf);
        assert!(!buf.is_empty());
    }

    #[test]
    fn test_avro_3433_recursive_definition_encode_array() {
        let mut buf = Vec::new();
        let schema = Schema::parse_str(
            r#"
        {
            "type":"record",
            "name":"TestStruct",
            "fields": [
                {
                    "name":"a",
                    "type":{
                        "type":"array",
                        "items": {
                            "type":"record",
                            "name": "Inner",
                            "fields": [ {
                                "name":"z",
                                "type":"int"
                            }]
                        }
                    }
                },
                {
                    "name":"b",
                    "type": {
                        "type":"map",
                        "values":"Inner"
                    }
                }
            ]
        }"#,
        )
        .unwrap();

        let inner_value1 = Value::Record(vec![("z".into(), Value::Int(3))]);
        let inner_value2 = Value::Record(vec![("z".into(), Value::Int(6))]);
        let outer_value = Value::Record(vec![
            ("a".into(), Value::Array(vec![inner_value1])),
            (
                "b".into(),
                Value::Map(vec![("akey".into(), inner_value2)].into_iter().collect()),
            ),
        ]);
        encode(&outer_value, &schema, &mut buf);
        assert!(!buf.is_empty());
    }

    #[test]
    fn test_avro_3433_recursive_definition_encode_map() {
        let mut buf = Vec::new();
        let schema = Schema::parse_str(
            r#"
        {
            "type":"record",
            "name":"TestStruct",
            "fields": [
                {
                    "name":"a",
                    "type":{
                        "type":"record",
                        "name": "Inner",
                        "fields": [ {
                            "name":"z",
                            "type":"int"
                        }]
                    }
                },
                {
                    "name":"b",
                    "type": {
                        "type":"map",
                        "values":"Inner"
                    }
                }
            ]
        }"#,
        )
        .unwrap();

        let inner_value1 = Value::Record(vec![("z".into(), Value::Int(3))]);
        let inner_value2 = Value::Record(vec![("z".into(), Value::Int(6))]);
        let outer_value = Value::Record(vec![
            ("a".into(), inner_value1),
            (
                "b".into(),
                Value::Map(vec![("akey".into(), inner_value2)].into_iter().collect()),
            ),
        ]);
        encode(&outer_value, &schema, &mut buf);
        assert!(!buf.is_empty());
    }

    #[test]
    fn test_avro_3433_recursive_definition_encode_record_wrapper() {
        let mut buf = Vec::new();
        let schema = Schema::parse_str(
            r#"
        {
            "type":"record",
            "name":"TestStruct",
            "fields": [
                {
                    "name":"a",
                    "type":{
                        "type":"record",
                        "name": "Inner",
                        "fields": [ {
                            "name":"z",
                            "type":"int"
                        }]
                    }
                },
                {
                    "name":"b",
                    "type": {
                        "type":"record",
                        "name": "InnerWrapper",
                        "fields": [ {
                            "name":"j",
                            "type":"Inner"
                        }]
                    }
                }
            ]
        }"#,
        )
        .unwrap();

        let inner_value1 = Value::Record(vec![("z".into(), Value::Int(3))]);
        let inner_value2 = Value::Record(vec![(
            "j".into(),
            Value::Record(vec![("z".into(), Value::Int(6))]),
        )]);
        let outer_value =
            Value::Record(vec![("a".into(), inner_value1), ("b".into(), inner_value2)]);
        encode(&outer_value, &schema, &mut buf);
        assert!(!buf.is_empty());
    }

    #[test]
    fn test_avro_3433_recursive_definition_encode_map_and_array() {
        let mut buf = Vec::new();
        let schema = Schema::parse_str(
            r#"
        {
            "type":"record",
            "name":"TestStruct",
            "fields": [
                {
                    "name":"a",
                    "type":{
                        "type":"map",
                        "values": {
                            "type":"record",
                            "name": "Inner",
                            "fields": [ {
                                "name":"z",
                                "type":"int"
                            }]
                        }
                    }
                },
                {
                    "name":"b",
                    "type": {
                        "type":"array",
                        "items":"Inner"
                    }
                }
            ]
        }"#,
        )
        .unwrap();

        let inner_value1 = Value::Record(vec![("z".into(), Value::Int(3))]);
        let inner_value2 = Value::Record(vec![("z".into(), Value::Int(6))]);
        let outer_value = Value::Record(vec![
            (
                "a".into(),
                Value::Map(vec![("akey".into(), inner_value2)].into_iter().collect()),
            ),
            ("b".into(), Value::Array(vec![inner_value1])),
        ]);
        encode(&outer_value, &schema, &mut buf);
        assert!(!buf.is_empty());
    }

    #[test]
    fn test_avro_3433_recursive_definition_encode_union() {
        let mut buf = Vec::new();
        let schema = Schema::parse_str(
            r#"
        {
            "type":"record",
            "name":"TestStruct",
            "fields": [
                {
                    "name":"a",
                    "type":["null", {
                        "type":"record",
                        "name": "Inner",
                        "fields": [ {
                            "name":"z",
                            "type":"int"
                        }]
                    }]
                },
                {
                    "name":"b",
                    "type":"Inner"
                }
            ]
        }"#,
        )
        .unwrap();

        let inner_value1 = Value::Record(vec![("z".into(), Value::Int(3))]);
        let inner_value2 = Value::Record(vec![("z".into(), Value::Int(6))]);
        let outer_value1 = Value::Record(vec![
            ("a".into(), Value::Union(1, Box::new(inner_value1))),
            ("b".into(), inner_value2.clone()),
        ]);
        encode(&outer_value1, &schema, &mut buf);
        assert!(!buf.is_empty());

        buf.drain(..);
        let outer_value2 = Value::Record(vec![
            ("a".into(), Value::Union(0, Box::new(Value::Null))),
            ("b".into(), inner_value2),
        ]);
        encode(&outer_value2, &schema, &mut buf);
        assert!(!buf.is_empty());
    }
}
